/*
    Dynamic Resolver
    5mukx
*/

use std::{
    ffi::c_void,
    ptr::null_mut,
};

use windows::{
    Win32::{
        Foundation::{FARPROC, HWND},
        System::{
            Diagnostics::Debug::{IMAGE_DATA_DIRECTORY, IMAGE_NT_HEADERS64},
            Kernel::LIST_ENTRY,
            SystemServices::{IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY},
        },
    },
    core::{PCSTR, s},
};

use windows::Win32::System::Threading::PEB;

use windows::Win32::System::WindowsProgramming::LDR_DATA_TABLE_ENTRY;

#[inline]
#[cfg(target_pointer_width = "64")]
pub fn __readgsqword(offset: u32) -> u64 {
    let out: u64;

    unsafe {
        std::arch::asm!(
            "mov {}, gs:[{:e}]",
            lateout(reg) out,
            in(reg) offset,
            options(nostack, pure, readonly),
        );
    }
    out
}

type MessageBoxA = unsafe extern "system" fn(HWND, PCSTR, PCSTR, u32) -> i32;
type LoadLibraryA = unsafe extern "system" fn(PCSTR) -> i32;

fn get_module_base_addr(module_name: &str) -> *mut c_void {
    unsafe {
        let peb_offset: *const u64 = __readgsqword(0x60) as *const u64;
        let peb: PEB = *(peb_offset as *const PEB);

        let mut p_ldr_data_table_entry: *const LDR_DATA_TABLE_ENTRY =
            (*peb.Ldr).InMemoryOrderModuleList.Flink as *const LDR_DATA_TABLE_ENTRY;
        let mut p_list_entry = &(*peb.Ldr).InMemoryOrderModuleList as *const LIST_ENTRY;

        loop {
            let buffer = std::slice::from_raw_parts(
                (*p_ldr_data_table_entry).FullDllName.Buffer.0,
                (*p_ldr_data_table_entry).FullDllName.Length as usize / 2,
            );
            let dll_name = String::from_utf16_lossy(buffer);

            if dll_name.to_lowercase().starts_with(module_name) {
                let module_base = (*p_ldr_data_table_entry).Reserved2[0];
                return module_base;
            }
            if p_list_entry == (*peb.Ldr).InMemoryOrderModuleList.Blink {
                println!("Module not found!");
                return null_mut();
            }
            p_list_entry = (*p_list_entry).Flink;
            p_ldr_data_table_entry = (*p_list_entry).Flink as *const LDR_DATA_TABLE_ENTRY;
        }
    }
}

fn get_proc_addr(module_handle: *mut c_void, function_name: &str) -> FARPROC {

    unsafe {
        let dos_headers = module_handle as *const IMAGE_DOS_HEADER;
    
        let nt_headers =
            (module_handle as u64 + (*dos_headers).e_lfanew as u64) as *const IMAGE_NT_HEADERS64;
    
        let data_directory =
            (&(*nt_headers).OptionalHeader.DataDirectory[0]) as *const IMAGE_DATA_DIRECTORY;
    
        let export_directory = (module_handle as u64 + (*data_directory).VirtualAddress as u64)
            as *const IMAGE_EXPORT_DIRECTORY;
    
        let mut address_array =
            (module_handle as u64 + (*export_directory).AddressOfFunctions as u64) as u64;
    
        let mut name_array =
            (module_handle as u64 + (*export_directory).AddressOfNames as u64) as u64;
        let mut name_ordinals =
            (module_handle as u64 + (*export_directory).AddressOfNameOrdinals as u64) as u64;

        loop {
            let name_offest: u32 = *(name_array as *const u32);
            let current_function_name =
                std::ffi::CStr::from_ptr((module_handle as u64 + name_offest as u64) as *const i8)
                    .to_str()
                    .unwrap();

            if current_function_name == function_name {
                address_array = address_array
                    + (*(name_ordinals as *const u16) as u64 * (std::mem::size_of::<u32>() as u64));
                let fun_addr: FARPROC = std::mem::transmute(
                    module_handle as u64 + *(address_array as *const u32) as u64,
                );
                return fun_addr;
            }

            name_array = name_array + std::mem::size_of::<u32>() as u64;
            name_ordinals = name_ordinals + std::mem::size_of::<u16>() as u64;
        }
    }
}

fn main() {
    unsafe {
        println!("[+] Getting base address of kernel32.dll");
        let kernel32_base_address = get_module_base_addr("kernel32.dll");

        println!("[+] Dynamically resolving LoadLibraryA");
        let dn_load_library_a: LoadLibraryA =
            std::mem::transmute(get_proc_addr(kernel32_base_address, "LoadLibraryA"));

        println!("[+] Load user32.dll");
        dn_load_library_a(s!("user32.dll"));

        println!("[+] Getting base address of user32.dll");
        let user32_base_address = get_module_base_addr("user32.dll");

        println!("[+] Dynamically resolved MessageBoxA");
        let dn_message_box_a: MessageBoxA =
            std::mem::transmute(get_proc_addr(user32_base_address, "MessageBoxA"));

        dn_message_box_a(
            HWND(null_mut()),
            s!("Hey, its Dynamically resolved"),
            s!("TEST"),
            0x0,
        );
    }
}
